#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import logging
import traceback
from store.store import BasicStore


# 存elasticsearch, es为 es连接, es_index为索引名
class EsStore(BasicStore):
    def __init__(self, config):
        self.es = config['es']
        self.es_index = config['es_index']

    def check_create(self, mapping):
        if self.es.indices.exists(index=self.es_index) is not True:
            res = self.es.indices.create(index=self.es_index, body=mapping)
            logging.info('es create index {0} res {1}'.format(self.es_index, res))

    def set(self, key, data, extra=None):
        query = {
            'query': {
                'term': {
                    '_id': key
                }
            }
        }
        try:
            search = self.es.search(index=self.es_index, body=query)
            # 无记录直接写入
            if len(search['hits']['hits']) == 0:
                result = self.es.index(index=self.es_index, id=key, body=data)
                if result['result'] == 'created':
                    return True
                else:
                    logging.warning('create res {0}'.format(result))
            else:
                # 有记录
                _source = search['hits']['hits'][0]['_source']
                if extra is None:
                    logging.warning('index {0} es_id {1} exists'.format(self.es_index, key))
                    return True
                else:
                    update_doc = None
                    if 'field' in extra:
                        field = extra['field']
                        # 检查字段是否存在，不存在则更新
                        if field is not None and field not in _source:
                            update_doc = {'doc': data}
                    elif 'update_doc' in extra:
                        # 检查是否有更新文档，有则更新
                        update_doc = {'doc': extra['update_doc']}
                    if update_doc is not None:
                        result = self.es.update(index=self.es_index, id=key, body=update_doc)
                        if result['result'] == 'updated':
                            return True
                        elif result['result'] == 'noop':
                            return True
                        else:
                            logging.warning('updated res {0}'.format(result))
        except Exception as e:
            logging.error(e)
            logging.error(traceback.format_exc())
        return False

    def get(self, key, extra=None):
        query = {
            'query': {
                'term': {
                    '_id': key
                }
            }
        }
        try:
            search = self.es.search(index=self.es_index, body=query)
            # 无记录
            if len(search['hits']['hits']) == 0:
                return None
            else:
                # 有记录, 返回第一个
                return search['hits']['hits'][0]['_source']
        except Exception as e:
            logging.error(e)
            logging.error(traceback.format_exc())
        return None
